{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import tensorflow as tf\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import lax\n",
    "from flax.nn.linear import _conv_dimension_numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 24, 24) [-0.2502736   0.34197763  0.81398004  0.8719954  -1.2945161   0.07460865\n",
      "  0.09205636 -0.24566422 -0.673575    0.24521644]\n",
      "(1, 1, 24, 24) [-0.2502736   0.3419779   0.81398     0.87199515 -1.2945158   0.07460855\n",
      "  0.09205629 -0.2456645  -0.673575    0.24521616]\n",
      "Error Rate: 0.006944444444444444\n"
     ]
    }
   ],
   "source": [
    "s = (16, 16) # strides\n",
    "k = 16\n",
    "x = np.random.randn(1, 3, 384, 384).astype(np.float32)\n",
    "\n",
    "th_x = torch.from_numpy(x)\n",
    "tf_x = np.float32(th_x.numpy()).transpose(0, 2, 3, 1)\n",
    "\n",
    "th_conv = torch.nn.Conv2d(3, 1, k, s, 0, bias=False)\n",
    "tf_conv = tf.keras.layers.Convolution2D(1, k, s, 'valid', use_bias=False)\n",
    "tf_conv(tf_x)\n",
    "tf_conv.set_weights([th_conv.weight.data.permute(2, 3, 1, 0).detach().numpy()])\n",
    "\n",
    "th_y = th_conv(th_x).detach().numpy()\n",
    "tf_y = tf_conv(tf_x).numpy()\n",
    "tf_y = tf_y.transpose(0, 3, 1, 2)\n",
    "\n",
    "print(th_y.shape, th_y.ravel()[:10])\n",
    "print(tf_y.shape, tf_y.ravel()[:10])\n",
    "print('Error Rate:', np.mean(np.abs(th_y - tf_y) > 1e-6))\n",
    "\n",
    "# assert np.allclose(th_y, tf_y, rtol=1e-5, atol=1e-8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 24, 24) [-0.2502736   0.34197763  0.81398004  0.8719954  -1.2945161   0.07460865\n",
      "  0.09205636 -0.24566422 -0.673575    0.24521644]\n",
      "(1, 1, 24, 24) [-0.2502736   0.3419779   0.81398     0.87199515 -1.2945158   0.07460855\n",
      "  0.09205629 -0.2456645  -0.673575    0.24521616]\n",
      "Error Rate: 0.006944444444444444\n"
     ]
    }
   ],
   "source": [
    "jax_x = jnp.asarray(th_x.numpy())\n",
    "jax_k = jnp.asarray(th_conv.weight.data.detach().numpy())\n",
    "\n",
    "jax_y = lax.conv_general_dilated(\n",
    "                 jax_x,    # lhs = NCHW image tensor\n",
    "                 jax_k, # rhs = OIHW conv kernel tensor\n",
    "                 s,  # window strides\n",
    "                 'VALID') # padding mode\n",
    "jax_y = jax_y._value\n",
    "\n",
    "print(th_y.shape, th_y.ravel()[:10])\n",
    "print(jax_y.shape, jax_y.ravel()[:10])\n",
    "print('Error Rate:', np.mean(np.abs(th_y - jax_y) > 1e-6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[[ 0.0000000e+00,  2.6822090e-07, -5.9604645e-08,\n",
       "          -2.3841858e-07,  2.3841858e-07, -9.6857548e-08,\n",
       "          -7.4505806e-08, -2.8312206e-07,  0.0000000e+00,\n",
       "          -2.8312206e-07,  3.5762787e-07,  2.3841858e-07,\n",
       "           4.7683716e-07, -2.9802322e-08,  1.7881393e-07,\n",
       "           1.1920929e-07, -2.3841858e-07, -1.0943040e-08,\n",
       "          -1.7881393e-07,  2.3841858e-07,  6.0535967e-09,\n",
       "           3.1292439e-07,  5.9604645e-08,  0.0000000e+00],\n",
       "         [ 2.3841858e-07,  2.9802322e-08,  7.4505806e-08,\n",
       "          -1.8253922e-07, -5.9604645e-08, -2.9802322e-08,\n",
       "          -3.2782555e-07,  1.7881393e-07, -1.4901161e-08,\n",
       "           7.4505806e-08, -1.0430813e-07,  1.3411045e-07,\n",
       "          -3.5762787e-07,  2.9802322e-08, -1.4901161e-07,\n",
       "           2.9802322e-08, -3.5762787e-07,  1.2665987e-07,\n",
       "           0.0000000e+00, -2.9802322e-08,  0.0000000e+00,\n",
       "           8.9406967e-08, -4.8428774e-08, -7.4505806e-08],\n",
       "         [ 0.0000000e+00,  3.5762787e-07, -2.7567148e-07,\n",
       "           5.9604645e-08,  8.9406967e-07,  0.0000000e+00,\n",
       "           2.0489097e-07, -1.7881393e-07,  3.7252903e-07,\n",
       "          -3.7252903e-07, -2.3841858e-07,  7.1525574e-07,\n",
       "          -2.3841858e-07, -2.3841858e-07, -7.1525574e-07,\n",
       "           1.1175871e-07, -6.7055225e-08,  3.5762787e-07,\n",
       "          -2.3841858e-07, -1.7881393e-07,  1.7881393e-07,\n",
       "          -6.5565109e-07,  3.2782555e-07,  6.5565109e-07],\n",
       "         [ 3.2782555e-07, -5.9604645e-08, -2.9802322e-08,\n",
       "          -1.2665987e-07, -1.7881393e-07,  2.6822090e-07,\n",
       "           1.7881393e-07, -2.9802322e-07, -2.2351742e-07,\n",
       "           3.5762787e-07,  1.8067658e-07,  5.9604645e-08,\n",
       "          -5.3644180e-07, -3.8743019e-07,  7.1525574e-07,\n",
       "           7.4505806e-07, -1.4901161e-07,  2.3841858e-07,\n",
       "           8.9406967e-08, -8.9406967e-08,  5.6624413e-07,\n",
       "          -8.3446503e-07,  3.5762787e-07,  1.7881393e-07],\n",
       "         [-1.3038516e-07, -1.7881393e-07,  2.9802322e-07,\n",
       "          -3.5762787e-07,  0.0000000e+00, -1.4901161e-07,\n",
       "          -7.4505806e-08, -8.9406967e-08,  2.6077032e-08,\n",
       "          -5.9604645e-07,  1.7881393e-07,  1.4901161e-07,\n",
       "           0.0000000e+00,  3.2782555e-07, -1.1920929e-07,\n",
       "          -8.1956387e-08,  2.9802322e-07, -5.9604645e-07,\n",
       "          -8.9406967e-08,  1.4901161e-07,  1.7881393e-07,\n",
       "           1.4901161e-07,  1.1920929e-07,  4.1723251e-07],\n",
       "         [-4.1723251e-07,  2.3841858e-07,  8.9406967e-08,\n",
       "          -2.1606684e-07, -3.5762787e-07, -1.7508864e-07,\n",
       "           3.3527613e-07,  0.0000000e+00, -2.6402995e-07,\n",
       "           0.0000000e+00, -2.3841858e-07, -4.7683716e-07,\n",
       "          -4.1723251e-07,  1.7881393e-07, -1.9371510e-07,\n",
       "           1.4901161e-07, -8.9406967e-08,  2.9802322e-08,\n",
       "           7.4505806e-09,  1.4901161e-08,  1.1920929e-07,\n",
       "           1.1920929e-07, -8.3446503e-07,  2.3841858e-07],\n",
       "         [ 1.1920929e-07, -2.9802322e-07,  1.1920929e-07,\n",
       "           3.5762787e-07,  1.7881393e-07,  8.9406967e-08,\n",
       "           1.4901161e-08,  3.5762787e-07,  4.7683716e-07,\n",
       "           7.4505806e-08,  1.5646219e-07, -5.9604645e-08,\n",
       "          -1.7881393e-07,  1.1920929e-07,  2.9802322e-07,\n",
       "           0.0000000e+00,  2.9802322e-08, -2.9802322e-07,\n",
       "          -4.7683716e-07, -4.4703484e-08,  4.4703484e-08,\n",
       "          -1.7881393e-07,  3.5762787e-07, -2.3841858e-07],\n",
       "         [-3.5762787e-07,  2.2351742e-07,  1.4901161e-07,\n",
       "           3.5762787e-07, -4.7683716e-07, -1.7881393e-07,\n",
       "           5.9604645e-08, -1.4901161e-07,  3.2782555e-07,\n",
       "           2.2351742e-07, -4.4703484e-08, -5.6624413e-07,\n",
       "           5.9604645e-08,  5.3644180e-07, -3.5762787e-07,\n",
       "           8.9406967e-08, -1.1175871e-07,  1.7881393e-07,\n",
       "          -2.3841858e-07, -2.0861626e-07,  2.7567148e-07,\n",
       "          -1.6391277e-07, -1.3411045e-07,  7.4505806e-08],\n",
       "         [-5.0663948e-07, -2.9802322e-08,  2.3841858e-07,\n",
       "           3.5762787e-07, -1.6391277e-07, -2.9802322e-07,\n",
       "           2.2351742e-08,  6.5565109e-07, -4.0233135e-07,\n",
       "           1.3783574e-07,  6.5565109e-07,  9.6857548e-08,\n",
       "          -1.3038516e-07, -1.4901161e-08, -2.9802322e-08,\n",
       "           4.7683716e-07,  1.7881393e-07,  5.9604645e-08,\n",
       "           8.9406967e-08, -1.7136335e-07, -1.1920929e-07,\n",
       "          -3.5762787e-07, -1.1175871e-08,  3.7252903e-09],\n",
       "         [-1.1920929e-07,  3.7252903e-08, -7.1525574e-07,\n",
       "           1.4901161e-07, -5.9604645e-08, -9.6857548e-08,\n",
       "          -1.0430813e-07,  9.6857548e-08, -5.9604645e-08,\n",
       "           1.7881393e-07, -1.7881393e-07, -2.2351742e-07,\n",
       "          -2.9802322e-07,  1.7881393e-07,  7.5902790e-08,\n",
       "           5.9604645e-08,  8.9406967e-08, -2.3841858e-07,\n",
       "          -5.9604645e-08, -5.9604645e-08,  1.6018748e-07,\n",
       "           2.6822090e-07,  3.5762787e-07,  0.0000000e+00],\n",
       "         [-1.4901161e-07,  0.0000000e+00,  0.0000000e+00,\n",
       "          -2.3841858e-07, -1.9371510e-07, -2.3841858e-07,\n",
       "          -2.3841858e-07,  2.6077032e-07,  8.9406967e-08,\n",
       "           1.4901161e-08, -2.9802322e-08,  5.9604645e-07,\n",
       "           2.0861626e-07,  9.5367432e-07, -6.2584877e-07,\n",
       "           1.1920929e-07, -5.9604645e-08, -1.1920929e-07,\n",
       "           7.4505806e-08,  1.1920929e-07,  2.9802322e-08,\n",
       "           3.8743019e-07,  5.9604645e-08,  2.6822090e-07],\n",
       "         [-5.2154064e-08,  8.9406967e-08,  2.9802322e-08,\n",
       "          -8.9406967e-08,  0.0000000e+00,  8.9406967e-08,\n",
       "           1.4901161e-07, -7.1525574e-07,  3.5762787e-07,\n",
       "           1.8626451e-07,  5.9604645e-08, -1.2665987e-07,\n",
       "          -5.9604645e-08, -8.9406967e-08,  5.9604645e-08,\n",
       "          -1.2479722e-07,  1.1920929e-06,  2.9802322e-07,\n",
       "           2.3841858e-07,  1.1920929e-07,  5.4389238e-07,\n",
       "           5.9604645e-08, -1.1920929e-07,  1.1920929e-07],\n",
       "         [ 0.0000000e+00, -1.1920929e-07, -5.9604645e-08,\n",
       "           8.9406967e-08, -5.2154064e-07, -7.3749106e-08,\n",
       "          -1.4901161e-07, -2.0861626e-07, -1.1920929e-07,\n",
       "           3.7252903e-08,  2.9802322e-07, -5.9604645e-08,\n",
       "          -2.9802322e-08,  3.2782555e-07, -5.9604645e-08,\n",
       "           8.9406967e-08, -1.1920929e-07,  4.3213367e-07,\n",
       "          -9.5367432e-07,  1.4901161e-07,  1.0430813e-07,\n",
       "           1.8626451e-08,  2.9802322e-07, -5.9604645e-08],\n",
       "         [ 3.5762787e-07,  1.4901161e-07, -6.3329935e-08,\n",
       "          -2.9802322e-07,  1.1920929e-07,  1.7881393e-07,\n",
       "           4.7683716e-07, -4.0046871e-08,  5.9604645e-08,\n",
       "           1.1920929e-07, -5.9604645e-08,  2.9802322e-07,\n",
       "          -1.4901161e-07, -3.0547380e-07, -2.0861626e-07,\n",
       "           5.9604645e-08, -5.9604645e-08,  1.1920929e-07,\n",
       "          -3.8743019e-07,  3.7252903e-08, -1.7881393e-07,\n",
       "          -1.7881393e-07,  0.0000000e+00,  1.3113022e-06],\n",
       "         [-2.9802322e-08,  0.0000000e+00,  9.6857548e-08,\n",
       "           8.9406967e-08,  2.3841858e-07, -2.9802322e-08,\n",
       "          -7.4505806e-08, -1.1920929e-07, -1.1920929e-07,\n",
       "          -5.9604645e-08,  6.7055225e-07, -2.3841858e-07,\n",
       "           4.1723251e-07, -2.9802322e-08, -8.9406967e-08,\n",
       "           2.9802322e-08, -1.3411045e-07,  8.9406967e-08,\n",
       "          -8.9406967e-08,  1.9371510e-07,  1.1920929e-07,\n",
       "          -1.5925616e-07,  2.9802322e-08, -4.8428774e-08],\n",
       "         [ 4.7683716e-07,  3.5762787e-07,  1.4901161e-07,\n",
       "           5.3644180e-07, -5.9604645e-08,  1.3113022e-06,\n",
       "           5.9604645e-08,  5.9604645e-08, -3.5762787e-07,\n",
       "           0.0000000e+00,  2.9802322e-07,  5.9604645e-08,\n",
       "          -5.9604645e-08,  0.0000000e+00, -6.5565109e-07,\n",
       "           1.9185245e-07, -1.1175871e-08, -1.1920929e-07,\n",
       "          -5.9604645e-08, -1.1920929e-07,  1.1920929e-07,\n",
       "           2.6822090e-07,  0.0000000e+00,  0.0000000e+00],\n",
       "         [ 1.7881393e-07,  1.4901161e-07, -5.9604645e-07,\n",
       "          -8.3446503e-07, -2.0861626e-07, -5.9604645e-07,\n",
       "          -3.5762787e-07, -5.9604645e-07,  0.0000000e+00,\n",
       "           6.7055225e-08, -1.1920929e-07, -3.6740676e-07,\n",
       "          -1.1920929e-07,  5.9604645e-08,  2.3841858e-07,\n",
       "           5.9604645e-07, -5.9604645e-08,  0.0000000e+00,\n",
       "          -1.3411045e-07,  1.1920929e-07,  0.0000000e+00,\n",
       "           5.9604645e-08,  1.6391277e-07, -5.9604645e-08],\n",
       "         [ 2.6822090e-07, -2.4586916e-07,  0.0000000e+00,\n",
       "           5.9604645e-08, -8.3446503e-07,  3.5762787e-07,\n",
       "          -4.7683716e-07,  2.5704503e-07,  0.0000000e+00,\n",
       "           8.9406967e-08, -6.5565109e-07, -2.3841858e-07,\n",
       "           2.6822090e-07, -5.9604645e-07,  7.4505806e-08,\n",
       "           6.5565109e-07, -1.7881393e-07,  3.8743019e-07,\n",
       "           2.3841858e-07, -8.9406967e-08, -3.3527613e-08,\n",
       "          -2.3841858e-07,  8.9406967e-08,  2.3841858e-07],\n",
       "         [ 3.2782555e-07,  2.8312206e-07,  1.7881393e-07,\n",
       "          -8.9406967e-08,  2.6822090e-07,  5.3085387e-08,\n",
       "           4.1723251e-07,  0.0000000e+00, -7.4505806e-08,\n",
       "          -2.9802322e-08,  1.1920929e-07, -2.0302832e-07,\n",
       "           2.3841858e-07,  1.6391277e-07, -2.9802322e-08,\n",
       "           4.7683716e-07, -1.7881393e-07,  5.9604645e-07,\n",
       "           1.3113022e-06,  0.0000000e+00,  4.4703484e-08,\n",
       "           3.5762787e-07,  2.9802322e-08, -8.9406967e-08],\n",
       "         [-4.7683716e-07, -5.9604645e-07, -3.2782555e-07,\n",
       "           1.7881393e-07,  1.7881393e-07, -3.5762787e-07,\n",
       "           2.0861626e-07,  4.4703484e-08,  2.9802322e-08,\n",
       "          -4.1723251e-07,  2.9802322e-08, -1.1920929e-07,\n",
       "          -2.5331974e-07,  0.0000000e+00,  0.0000000e+00,\n",
       "           0.0000000e+00, -1.1920929e-07,  2.9802322e-08,\n",
       "          -4.4703484e-07,  2.9802322e-07,  4.7683716e-07,\n",
       "          -2.9802322e-07, -2.9802322e-08,  0.0000000e+00],\n",
       "         [-5.4016709e-08, -4.1723251e-07,  2.2351742e-08,\n",
       "           7.4505806e-08,  5.9604645e-08, -5.9604645e-08,\n",
       "           2.9802322e-08,  0.0000000e+00, -1.4901161e-07,\n",
       "           2.9802322e-08, -3.5762787e-07,  3.5762787e-07,\n",
       "          -2.7194619e-07,  5.9604645e-08, -5.9604645e-08,\n",
       "           7.1525574e-07,  5.9604645e-08,  7.1525574e-07,\n",
       "          -5.9604645e-08,  4.1723251e-07, -7.4505806e-08,\n",
       "           2.3841858e-07, -2.9802322e-08,  4.4703484e-07],\n",
       "         [ 8.3446503e-07,  3.2782555e-07,  3.5762787e-07,\n",
       "          -1.1920929e-07,  7.4505806e-08,  1.7881393e-07,\n",
       "           0.0000000e+00, -5.9604645e-07,  2.9802322e-07,\n",
       "          -1.1920929e-07, -2.3841858e-07, -7.4505806e-08,\n",
       "           2.9802322e-08,  5.9604645e-08, -6.7055225e-08,\n",
       "           1.1920929e-07, -1.4901161e-07, -1.9371510e-07,\n",
       "          -5.9604645e-07, -1.3411045e-07, -7.4505806e-08,\n",
       "          -2.2351742e-07,  1.7881393e-07,  8.9406967e-08],\n",
       "         [ 2.3841858e-07,  0.0000000e+00, -5.9604645e-08,\n",
       "           1.1920929e-07,  1.1920929e-07,  2.0861626e-07,\n",
       "          -2.9802322e-07,  1.7695129e-08, -2.3841858e-07,\n",
       "           2.9802322e-07, -3.5762787e-07, -5.9604645e-08,\n",
       "          -8.9406967e-08,  0.0000000e+00, -1.4901161e-07,\n",
       "           5.9604645e-08,  2.0861626e-07, -5.9604645e-08,\n",
       "           2.9802322e-08, -1.4901161e-08,  5.9604645e-08,\n",
       "           0.0000000e+00, -2.3841858e-07,  4.7683716e-07],\n",
       "         [ 5.9604645e-08, -2.0861626e-07, -2.2351742e-07,\n",
       "          -8.9406967e-08, -1.1920929e-07,  1.4156103e-07,\n",
       "           0.0000000e+00,  4.1723251e-07,  2.0861626e-07,\n",
       "          -8.1956387e-08,  1.1920929e-07, -2.9802322e-08,\n",
       "          -1.7881393e-07,  5.9604645e-08,  5.9604645e-07,\n",
       "          -7.1525574e-07, -3.7252903e-07, -8.1956387e-08,\n",
       "          -4.7683716e-07, -2.0861626e-07, -7.4505806e-08,\n",
       "          -8.3819032e-08, -7.1525574e-07, -1.1920929e-07]]]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax_y - th_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 24, 24) [-0.2502736   0.34197763  0.81398004  0.8719954  -1.2945161   0.07460865\n",
      "  0.09205636 -0.24566422 -0.673575    0.24521644]\n",
      "(1, 1, 24, 24) [-0.2502736   0.3419779   0.81398     0.87199515 -1.2945158   0.07460855\n",
      "  0.09205629 -0.2456645  -0.673575    0.24521616]\n",
      "Error Rate: 0.006944444444444444\n"
     ]
    }
   ],
   "source": [
    "jax_x = jnp.asarray(th_x.numpy().transpose(0, 2, 3, 1))\n",
    "jax_k = jnp.asarray(th_conv.weight.data.detach().numpy().transpose(2, 3, 1, 0))\n",
    "\n",
    "dimension_numbers = _conv_dimension_numbers(jax_x.shape)\n",
    "jax_y = lax.conv_general_dilated(\n",
    "                 jax_x,    # lhs = NHWc image tensor\n",
    "                 jax_k,    # rhs = HWIO conv kernel tensor\n",
    "                 s,        # window strides\n",
    "                 'VALID',  # padding mode\n",
    "                 dimension_numbers=dimension_numbers,)  \n",
    "jax_y = jax_y._value.transpose(0, 3, 1, 2)\n",
    "\n",
    "print(th_y.shape, th_y.ravel()[:10])\n",
    "print(jax_y.shape, jax_y.ravel()[:10])\n",
    "print('Error Rate:', np.mean(np.abs(th_y - jax_y) > 1e-6))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
